import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression

pd.set_option("mode.chained_assignment", None)
import warnings

warnings.filterwarnings("ignore")
import akshare as ak

SPX_500_index = ak.index_us_stock_sina(symbol=".INX")
Hang_Seng_index = ak.stock_hk_daily(symbol="HSI")
CSI_300_index = ak.stock_zh_index_daily(symbol="sh000300")

# print(SPX_500_index)


def RSRS(data, N, S1, S2):
    #     for df in data.rolling(N):
    #         print(df)
    # Calculate daily increase based on closing price
    data["pct"] = data["close"] / data["close"].shift(1) - 1.0

    # Calculate slope using the least squares method
    def calculate_beta(df, window=N):
        if df.shape[0] < window:
            return np.nan
        x = df["low"].values
        y = df["high"].values
        beta = LinearRegression().fit(x.reshape(-1, 1), y).coef_[0]
        return beta

    #     print(data.rolling(N))
    #     print(type(data.rolling(N)))
    data["beta"] = [calculate_beta(df, window=N) for df in data.rolling(N)]
    data1 = data.dropna().copy().reset_index(drop=True)

    # Initialize buying and selling tags and positions
    data1["flag"] = 0
    data1["position"] = 0
    position = 0

    # Construct a trading strategy based on the relationship between slope and threshold
    for i in range(1, data1.shape[0] - 1):
        beta = data1.loc[i, "beta"]
        if (position == 0) and (beta > S1):
            data1.loc[i, "flag"] = 1
            data1.loc[i + 1, "position"] = 1
            position = 1
        elif (position == 1) and (beta < S2):
            data1.loc[i, "flag"] = -1
            data1.loc[i + 1, "position"] = 0
            position = 0
        else:
            data1.loc[i + 1, "position"] = data1.loc[i, "position"]

    # Calculate the net value of the strategy and draw a graph to display it
    data1["strategy_pct"] = data1["pct"] * data1["position"]
    # print(data1['strategy_pct'])
    data1["strategy"] = (1.0 + data1["strategy_pct"]).cumprod()
    data1["Hang_Seng_index"] = (1.0 + data1["pct"]).cumprod()
    annual_return = 100 * (pow(data1["strategy"].iloc[-1], 250 / data1.shape[0]) - 1.0)
    data1.index = pd.to_datetime(data1["date"])
    ax = data1[["strategy", "Hang_Seng_index"]].plot(
        figsize=(16, 8), color=["SteelBlue", "Red"], title="RSRS"
    )
    plt.show()
    print(data1["beta"].mean(), data1["beta"].std())
    print(data1["beta"])
    data1.to_csv("df.csv")


RSRS(Hang_Seng_index, N=15, S1=1.0, S2=0.8)
